from transformers import AutoModel
import onnx
import torch

# 加载模型
model = AutoModel.from_pretrained("jackietung/bert-base-chinese-finetuned-sentiment")
# 构造 dummy 输入（根据模型输入要求调整，比如 BERT 通常是 input_ids, attention_mask 等）
dummy_input = torch.randint(0, 1000, (1, 128))  # 假设输入序列长度 128，batch 为 1
# 转换为 onnx 格式，保存到指定路径
torch.onnx.export(model, dummy_input, "models/bert_chinese_sentiment.onnx",
                  input_names=["input_ids"],  # 输入名称，按需改
                  output_names=["output"],  # 输出名称，按需改
                  dynamic_axes={"input_ids": {0: "batch_size", 1: "seq_length"}})  # 动态维度（如果需要）